# coding: utf-8
from __future__ import print_function
from __future__ import division
import tensorflow as tf
from nets import nets_factory
from preprocessing import preprocessing_factory
import reader
import model
import time
import losses
import utils
import os
import argparse

slim = tf.contrib.slim


def parse_args():#定义域名空间，训练的配置文件.yml
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--conf', default='conf/mosaic.yml', help='the path to the conf file')
    return parser.parse_args()


def main(FLAGS):
    # 获取style image的style feature
    style_features_t = losses.get_style_features(FLAGS)#获取风格的特征

    # Make sure the training path exists.
    training_path = os.path.join(FLAGS.model_path, FLAGS.naming)

    if not(os.path.exists(training_path)):
    #创建training path
        os.makedirs(training_path)

    with tf.Graph().as_default():#创建默认的图
        with tf.Session() as sess:
            """Build Network建立训练网络"""
            # network_fn是损失网络函数，不需要对损失网络进行训练
            network_fn = nets_factory.get_network_fn(#返回一个网络 损失模型 VGG16
                FLAGS.loss_model,
                num_classes=1,
                is_training=False)
            #损失网络中要用的图像的预处理函数
            """loss_model=vgg16"""
            image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
                FLAGS.loss_model,
                is_training=False)
            '''读取数据集，返回处理过的图片张量'''
            processed_images = reader.image(FLAGS.batch_size, FLAGS.image_size, FLAGS.image_size,
                                            'train2014/', image_preprocessing_fn, epochs=FLAGS.epoch)
            ''' # 使用图像生成网络，输入原始图像，输出生成图像'''
            generated = model.net(processed_images, training=True)
            '''unstack()矩阵的分解
            因为generated同样需要送到损失网络中计算loss，所以需要image_preprocessing_fn进行处理'''
            processed_generated = [image_preprocessing_fn(image, FLAGS.image_size, FLAGS.image_size)
                                   for image in tf.unstack(generated, axis=0, num=FLAGS.batch_size)
                                   ]
            ''' 把图片list组成一个batch'''
            #tf.stack矩阵在0维度拼接，将一组R维张量变为R+1维张量
            processed_generated = tf.stack(processed_generated)
            '''tf.concat()在0维上连接
             将原始图像，生成图像送到损失网络中
               network_fn（）返回一个vgg16网络net和一个字典
            '''
            _, endpoints_dict = network_fn(tf.concat([processed_generated, processed_images], 0), spatial_squeeze=False)

            # Log the structure of loss network
            tf.logging.info('Loss network layers(You can define them in "content_layers" and "style_layers"):')
            for key in endpoints_dict:
                tf.logging.info(key)

            """Build Losses"""
            content_loss = losses.content_loss(endpoints_dict, FLAGS.content_layers)
            style_loss, style_loss_summary = losses.style_loss(endpoints_dict, style_features_t, FLAGS.style_layers)
            tv_loss = losses.total_variation_loss(generated)  # use the unprocessed image
            '''最终的损失=权重*损失'''
            loss = FLAGS.style_weight * style_loss + FLAGS.content_weight * content_loss + FLAGS.tv_weight * tv_loss

            # Add Summary for visualization in tensorboard.在TensorBoard中添加可视化摘要。
            """Add Summary"""
            tf.summary.scalar('losses/content_loss', content_loss)
            tf.summary.scalar('losses/style_loss', style_loss)
            tf.summary.scalar('losses/regularizer_loss', tv_loss)

            tf.summary.scalar('weighted_losses/weighted_content_loss', content_loss * FLAGS.content_weight)
            tf.summary.scalar('weighted_losses/weighted_style_loss', style_loss * FLAGS.style_weight)
            tf.summary.scalar('weighted_losses/weighted_regularizer_loss', tv_loss * FLAGS.tv_weight)
            tf.summary.scalar('total_loss', loss)

            for layer in FLAGS.style_layers:
                tf.summary.scalar('style_losses/' + layer, style_loss_summary[layer])
            tf.summary.image('generated', generated)
            # tf.image_summary('processed_generated', processed_generated)  # May be better?
            tf.summary.image('origin', tf.stack([
                image_unprocessing_fn(image) for image in tf.unstack(processed_images, axis=0, num=FLAGS.batch_size)
            ]))
            summary = tf.summary.merge_all()
            writer = tf.summary.FileWriter(training_path)#文件写入

            """Prepare to Train
                定义全局step,初始为0
            """
            global_step = tf.Variable(0, name="global_step", trainable=False)
            # 找出需要训练的变量，本项目中，只需要训练图像生成网络中的变量，不需要训练损失网络中的变量
            variable_to_train = []
            # 使用tf.trainable_variables()找出所有可以训练的变量
            for variable in tf.trainable_variables():
                # 如果不在损失网络中，把他们加入列表variables_to_train
                if not(variable.name.startswith(FLAGS.loss_model)):
                    variable_to_train.append(variable)
            # Adam优化算法：是一个寻找全局最优点的优化算法，引入了二次方梯度校正。
            # 相比于基础SGD算法，1.不容易陷于局部优点。2.速度更快
            train_op = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step, var_list=variable_to_train)

            variables_to_restore = []
            for v in tf.global_variables():
                if not(v.name.startswith(FLAGS.loss_model)):
                    variables_to_restore.append(v)
            #保存检查点
            saver = tf.train.Saver(variables_to_restore, write_version=tf.train.SaverDef.V1)

            sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])

            # Restore variables for loss network.恢复损失网络的变量
            # 载入损失网络初始权重
            init_func = utils._get_init_fn(FLAGS)
            init_func(sess)

            # Restore variables for training model if the checkpoint file exists.
            '''如果检查点文件存在，则还原训练模型的变量'''
            last_file = tf.train.latest_checkpoint(training_path)
            if last_file:
                tf.logging.info('Restoring model from {}'.format(last_file))
                saver.restore(sess, last_file)

            """Start Training"""
            coord = tf.train.Coordinator()#使用 tf.train.Coordinator()来创建一个线程管理器（协调器）对象。
            threads = tf.train.start_queue_runners(coord=coord)#使用start_queue_runners 启动队列填充
            start_time = time.time()
            try:
                while not coord.should_stop():#判断coord.request_stop()函数
                    _, loss_t, step = sess.run([train_op, loss, global_step])
                    elapsed_time = time.time() - start_time
                    start_time = time.time()
                    """logging"""
                    # print(step)
                    if step % 10 == 0:
                        tf.logging.info('step: %d,  total Loss %f, secs/step: %f' % (step, loss_t, elapsed_time))
                    """summary"""
                    if step % 25 == 0:
                        tf.logging.info('adding summary...')
                        summary_str = sess.run(summary)
                        writer.add_summary(summary_str, step)
                        writer.flush()
                    """checkpoint"""
                    if step % 1000 == 0:
                        saver.save(sess, os.path.join(training_path, 'fast-style-model.ckpt'), global_step=step)
            except tf.errors.OutOfRangeError:
                saver.save(sess, os.path.join(training_path, 'fast-style-model.ckpt-done'))
                tf.logging.info('Done training -- epoch limit reached')
            finally:
                coord.request_stop()
            coord.join(threads)#把开启的线程加入主线程，等待threads结束


if __name__ == '__main__':
    tf.logging.set_verbosity(tf.logging.INFO)# 设定日志等级
    args = parse_args()
    FLAGS = utils.read_conf_file(args.conf)#读取训练的配置文件
    main(FLAGS)
